# -*- coding: utf-8 -*-
# time: 2025/5/15 09:06
# file: ch01.py
# author: hanson
"""
sk-62fb03f9198f4720bb359c873757fafc
https://bailian.console.aliyun.com/?tab=app#/api-key
"""
from langchain.agents import AgentExecutor, create_tool_calling_agent, initialize_agent, AgentType
from langchain.tools import Tool
from langchain_community.chat_models.tongyi import ChatTongyi
from langchain_community.utilities.sql_database import SQLDatabase
from langchain_core.prompts import ChatPromptTemplate
from langchain_ollama import ChatOllama

# 1. 初始化Qwen大模型 (通过Dashscope) Qwen模型另外处理
# llm = ChatTongyi(
#     model_name="qwen-turbo",  # qwen-turbo,qwen-max
#     temperature =0.3,
#     dashscope_api_key="sk-62fb03f9198f4720bb359c873757fafc",  # 替换为实际API key
# )
llm = ChatOllama(model="llama3.2:3b", temperature=0.1)
# 2. 连接MySQL数据库
db = SQLDatabase.from_uri(
    "mysql+pymysql://root:1234qwer@192.168.6.43:3306/iot_ts",
    sample_rows_in_table_info=3
)

# 3. 创建数据库查询工具
def query_database(query: str) -> str:
    """执行SQL查询并返回结果"""
    try:
        result = db.run(query)
        return str(result)
    except Exception as e:
        return f"查询错误: {str(e)}"


db_tool = Tool(
    name="database_query",
    func=query_database,
    description="用于查询MySQL数据库的工具。输入应该是有效的SQL查询语句。"
)


# 4. 创建自定义加法工具
def add_numbers(a: float, b: float) -> str:
    """将两个数字相加"""
    print(f"日志==》正在计算：{a} + {b}")
    return str(a + b)


add_tool = Tool(
    name="add_numbers",
    func=add_numbers,
    description="用于将两个数字相加的工具。输入应该是两个数字，例如 '3, 5'。"
)

# 5. 构建Agent
tools = [db_tool, add_tool]

prompt = ChatPromptTemplate.from_messages([
    ("system", "你是专业的AI助手，可以访问数据库和执行计算。"),
    ("human", "{input}"),
    ("placeholder", "{agent_scratchpad}")
])

agent = create_tool_calling_agent(llm, tools, prompt)
# 新版推荐使用ZERO_SHOT_REACT_DESCRIPTION类型
# agent = initialize_agent(
#     tools,
#     llm,
#     agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
#     verbose=True,
#     handle_parsing_errors=True  # 关键：处理解析错误
# )

agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

# 6. 测试运行
if __name__ == "__main__":
    # 数据库查询示例
    print("=== 数据库查询 ===")
    print(agent_executor.invoke({
        "input": "查询users表的前5条记录"
    }))

    # 加法计算示例
    print("\n=== 加法计算 ===")
    print(agent_executor.invoke({
        "input": "计算123加456等于多少"
    }))

    # 混合任务示例
    print("\n=== 混合任务 ===")
    print(agent_executor.invoke({
        "input": "先计算100加200，然后查询tk_device表有多少条记录"
    }))
